/*
给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。

 

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
 

提示:

1 <= nums1.length, nums2.length <= 104
-109 <= nums1[i], nums2[i] <= 109
nums1, nums2 均为升序排列
1 <= k <= 1000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
*/

#include "../stdc++.h"

// 优先级队列
// 时间复杂度：O(KlogK)
// 空间复杂度：O(K)
class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        auto cmp = [&nums1, &nums2] (const pair<int, int>& a, const pair<int, int>& b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };
        vector<vector<int>> res;
        int m = nums1.size();
        int n = nums2.size();
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> maxHeap(cmp);
        for (int i{0}; i < min(k, m); ++i) {
            maxHeap.emplace(i, 0);
        }
        while (k-- > 0 && !maxHeap.empty()) {
            int x = maxHeap.top().first;
            int y = maxHeap.top().second;
            maxHeap.pop();
            res.emplace_back(initializer_list<int>{nums1[x], nums2[y]});
            if (y + 1 < n) {
                maxHeap.emplace(x, y + 1);
            }
        }
        return res;
    }
};

// 二分查找
class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int m = nums1.size();
        int n = nums2.size();
        auto count = [&] (int target) {
            long long cnt = 0;
            int start = 0;
            int end = n - 1;
            while (start < m && end >= 0) {
                if (nums1[start] + nums2[end] > target) {
                    --end;
                } else {
                    cnt += end + 1;
                    ++start;
                }
            }
            return cnt;
        };

        // 二分查找第 k 小的数对和的大小
        int left = nums1[0] + nums2[0];
        int right = nums1.back() + nums2.back();
        int pairSum = right;
        while (left <= right) {
            int mid = left + ((right - left) >> 1);
            if (count(mid) < k) {
                left = mid + 1;
            } else {
                pairSum = mid;
                right = mid - 1;
            }
        }

        vector<vector<int>> res;
        int pos = n - 1;
        // 找到小于目标值 pairSum 的数对
        for (int i = 0; i < m; ++i) {
            while (pos >= 0 && nums1[i] + nums2[pos] >= pairSum) {
                --pos;
            }
            for (int j = 0; j <= pos && k > 0; ++j, --k) {
                res.push_back({nums1[i], nums2[j]});
            }
        }
        // 找到等于目标值 pairSum 的数对
        pos = n - 1;
        for (int i = 0; i < m && k > 0; i++) {
            while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {
                pos--;
            }
            for (int j = i; k > 0 && j >= 0 && nums1[j] + nums2[pos] == pairSum; j--, k--) {
                res.push_back({nums1[i], nums2[pos]});
            }
        }
        return res;
    }
};
